from torch import nn
from efficientnet_pytorch import EfficientNet


class Efficientnet(nn.Module):
    def __init__(self, num_classes=2 ,model_name = 'efficientnet-b0'):
        super(Efficientnet, self).__init__()
        model = EfficientNet.from_pretrained(model_name, advprop=True)
        # model = EfficientNet.from_pretrained('efficientnet-b7')
        in_channel = model._fc.in_features
        model._fc = nn.Linear(in_channel, num_classes)
        self.efficientnet = model

    def forward(self, img):
        out = self.efficientnet(img)
        return out